from DOA import DOA
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm

def run_simulation(doa, doa_methods, fixed_angles, snr_range, snr_steps, num_repeats=50):
    num_methods = len(doa_methods)
    snr_values = np.linspace(snr_range[0], snr_range[1], snr_steps)
    estimated_angles = np.zeros((num_methods, snr_steps, doa.num_signals))  # 存储估计角度

    for snr_idx, snr in enumerate(tqdm(snr_values)):
        # 针对每个 SNR 重复多次实验
        count_success = [0] * num_methods
        for _ in range(num_repeats):
            signal, _ = doa.generate_signal(SNR=snr, print_angle=False, fixed_angle=fixed_angles)
            for method_idx, method in enumerate(doa_methods):
                estimated_theta = method(signal)
                # doa.draw_spectrum_result(doa.spectrum, estimated_theta)
                # 检查是否估计到角度
                if estimated_theta.size == doa.num_signals:
                    count_success[method_idx] += 1
                    estimated_angles[method_idx, snr_idx] += estimated_theta  # 累加角度并求平均
                else:
                    estimated_angles[method_idx, snr_idx] = 0  # 如果没有估计到角度，则设置为 0

        for method_idx in range(num_methods):
            if count_success[method_idx] > 0:
                estimated_angles[method_idx, snr_idx] /= count_success[method_idx]
            else:
                estimated_angles[method_idx, snr_idx] = 0  # 如果没有成功实验，则设置为 0

    # 绘制结果
    plt.figure()
    colors = ['r', 'g', 'b', 'm']  # 为不同方法设置不同颜色
    method_names = [method.__name__ for method in doa_methods]
    for method_idx in range(num_methods):
        for signal_idx in range(doa.num_signals):
            plt.plot(snr_values, np.rad2deg(estimated_angles[method_idx, :, signal_idx]), color=colors[method_idx], linestyle='-', label=f'{method_names[method_idx]} (Signal {signal_idx+1})')

    plt.title('DOA Estimation vs SNR (Fixed Angles, 100 Repeats)')
    plt.xlabel('SNR (dB)')
    plt.ylabel('Estimated Angle (degrees)')
    plt.legend()
    plt.grid(True)
    plt.show()

if __name__ == "__main__":
    doa_methods = (DOA().music, DOA().root_music, DOA().esprit, DOA().TLS_ESPRIT)
    import time
    doa = DOA()
    # doa_methods = (doa.music, )
    fixed_angles = [np.deg2rad(20), np.deg2rad(110)]  # 设置固定角度
    snr_range = (-10, 10)  # SNR 范围
    snr_steps = 20  # SNR 步数
    start_time = time.time()
    run_simulation(doa, doa_methods, fixed_angles, snr_range, snr_steps, num_repeats=50)
    end_time = time.time()
    print(f"程序执行时间为：{end_time - start_time} 秒")

